"""WebSocket处理器 - 处理WebSocket的MCP通信"""

import asyncio
import json
from typing import Dict, Any, Optional, Set
from loguru import logger

try:
    import websockets
    from websockets.server import WebSocketServerProtocol
    from websockets.exceptions import ConnectionClosed, WebSocketException
except ImportError:
    logger.warning("websockets库未安装，WebSocket功能将不可用")
    websockets = None
    WebSocketServerProtocol = None
    ConnectionClosed = Exception
    WebSocketException = Exception

from .server import MCPServer


class WebSocketHandler:
    """WebSocket处理器"""
    
    def __init__(self, server: MCPServer, config: Dict[str, Any]):
        if websockets is None:
            raise ImportError("websockets库未安装，请运行: pip install websockets")
        
        self.server = server
        self.config = config
        self.ws_config = config.get('mcp', {}).get('websocket', {})
        
        # WebSocket配置
        self.host = self.ws_config.get('host', 'localhost')
        self.port = self.ws_config.get('port', 8765)
        self.ping_interval = self.ws_config.get('ping_interval', 20)
        self.ping_timeout = self.ws_config.get('ping_timeout', 10)
        self.max_size = self.ws_config.get('max_size', 1024 * 1024)  # 1MB
        self.max_queue = self.ws_config.get('max_queue', 32)
        
        # 运行状态
        self.is_running = False
        self.websocket_server = None
        self.connections: Dict[str, WebSocketServerProtocol] = {}
        
        logger.info(f"WebSocket处理器已初始化 - {self.host}:{self.port}")
    
    async def start(self):
        """启动WebSocket处理器"""
        if self.is_running:
            logger.warning("WebSocket处理器已在运行")
            return
        
        try:
            # 启动MCP服务器
            await self.server.start()
            
            # 启动WebSocket服务器
            self.websocket_server = await websockets.serve(
                self._handle_connection,
                self.host,
                self.port,
                ping_interval=self.ping_interval,
                ping_timeout=self.ping_timeout,
                max_size=self.max_size,
                max_queue=self.max_queue
            )
            
            self.is_running = True
            logger.info(f"WebSocket服务器已启动 - ws://{self.host}:{self.port}")
            
        except Exception as e:
            logger.error(f"启动WebSocket处理器失败: {e}")
            raise
    
    async def stop(self):
        """停止WebSocket处理器"""
        if not self.is_running:
            logger.warning("WebSocket处理器未在运行")
            return
        
        try:
            self.is_running = False
            
            # 关闭所有WebSocket连接
            if self.connections:
                close_tasks = []
                for connection_id, websocket in list(self.connections.items()):
                    close_tasks.append(self._close_connection(connection_id, websocket))
                
                if close_tasks:
                    await asyncio.gather(*close_tasks, return_exceptions=True)
            
            # 停止WebSocket服务器
            if self.websocket_server:
                self.websocket_server.close()
                await self.websocket_server.wait_closed()
                self.websocket_server = None
            
            # 停止MCP服务器
            await self.server.stop()
            
            logger.info("WebSocket处理器已停止")
            
        except Exception as e:
            logger.error(f"停止WebSocket处理器失败: {e}")
    
    async def _handle_connection(self, websocket: WebSocketServerProtocol, path: str):
        """处理WebSocket连接"""
        connection_id = f"ws_{id(websocket)}"
        
        try:
            # 记录连接信息
            client_info = {
                'transport': 'websocket',
                'remote_address': websocket.remote_address,
                'path': path,
                'subprotocols': websocket.subprotocols,
                'headers': dict(websocket.request_headers) if hasattr(websocket, 'request_headers') else {}
            }
            
            # 连接到MCP服务器
            await self.server.connect_client(connection_id, client_info)
            self.connections[connection_id] = websocket
            
            logger.info(f"WebSocket客户端已连接: {connection_id} from {websocket.remote_address}")
            
            # 处理消息循环
            await self._message_loop(websocket, connection_id)
            
        except ConnectionClosed:
            logger.info(f"WebSocket连接已关闭: {connection_id}")
        except WebSocketException as e:
            logger.error(f"WebSocket异常: {connection_id} - {e}")
        except Exception as e:
            logger.error(f"处理WebSocket连接失败: {connection_id} - {e}")
        finally:
            # 清理连接
            await self._cleanup_connection(connection_id)
    
    async def _message_loop(self, websocket: WebSocketServerProtocol, connection_id: str):
        """消息处理循环"""
        try:
            async for message in websocket:
                try:
                    # 检查消息类型
                    if isinstance(message, bytes):
                        message = message.decode('utf-8')
                    
                    logger.debug(f"收到WebSocket消息: {connection_id} - {len(message)} 字符")
                    
                    # 处理消息
                    response = await self.server.handle_message(message, connection_id)
                    
                    # 发送响应
                    await websocket.send(response)
                    
                    logger.debug(f"发送WebSocket响应: {connection_id} - {len(response)} 字符")
                    
                except json.JSONDecodeError as e:
                    logger.error(f"JSON解析失败: {connection_id} - {e}")
                    # 发送错误响应
                    error_response = {
                        'jsonrpc': '2.0',
                        'error': {
                            'code': -32700,
                            'message': 'Parse error',
                            'data': str(e)
                        },
                        'id': None
                    }
                    await websocket.send(json.dumps(error_response, ensure_ascii=False))
                
                except Exception as e:
                    logger.error(f"处理WebSocket消息失败: {connection_id} - {e}")
                    # 发送内部错误响应
                    error_response = {
                        'jsonrpc': '2.0',
                        'error': {
                            'code': -32603,
                            'message': 'Internal error',
                            'data': str(e)
                        },
                        'id': None
                    }
                    await websocket.send(json.dumps(error_response, ensure_ascii=False))
        
        except ConnectionClosed:
            logger.debug(f"WebSocket连接已关闭: {connection_id}")
        except Exception as e:
            logger.error(f"WebSocket消息循环异常: {connection_id} - {e}")
    
    async def _cleanup_connection(self, connection_id: str):
        """清理连接"""
        try:
            # 从连接字典中移除
            if connection_id in self.connections:
                del self.connections[connection_id]
            
            # 从MCP服务器断开
            await self.server.disconnect_client(connection_id)
            
            logger.debug(f"WebSocket连接已清理: {connection_id}")
            
        except Exception as e:
            logger.error(f"清理WebSocket连接失败: {connection_id} - {e}")
    
    async def _close_connection(self, connection_id: str, websocket: WebSocketServerProtocol):
        """关闭单个连接"""
        try:
            if not websocket.closed:
                await websocket.close()
            await self._cleanup_connection(connection_id)
        except Exception as e:
            logger.error(f"关闭WebSocket连接失败: {connection_id} - {e}")
    
    async def broadcast_message(self, message: str, exclude_connections: Optional[Set[str]] = None):
        """向所有连接广播消息"""
        if not self.connections:
            logger.debug("没有WebSocket连接可广播")
            return
        
        exclude_connections = exclude_connections or set()
        broadcast_tasks = []
        
        for connection_id, websocket in self.connections.items():
            if connection_id not in exclude_connections and not websocket.closed:
                broadcast_tasks.append(self._send_to_connection(websocket, message, connection_id))
        
        if broadcast_tasks:
            results = await asyncio.gather(*broadcast_tasks, return_exceptions=True)
            
            # 统计发送结果
            success_count = sum(1 for result in results if not isinstance(result, Exception))
            error_count = len(results) - success_count
            
            logger.debug(f"广播消息完成: 成功 {success_count}, 失败 {error_count}")
    
    async def _send_to_connection(self, websocket: WebSocketServerProtocol, message: str, connection_id: str):
        """向单个连接发送消息"""
        try:
            await websocket.send(message)
            return True
        except ConnectionClosed:
            logger.debug(f"连接已关闭，无法发送消息: {connection_id}")
            # 清理已关闭的连接
            await self._cleanup_connection(connection_id)
            return False
        except Exception as e:
            logger.error(f"发送消息到WebSocket连接失败: {connection_id} - {e}")
            return False
    
    async def send_notification(self, method: str, params: Dict[str, Any], 
                              target_connections: Optional[Set[str]] = None):
        """发送通知到指定连接或所有连接"""
        try:
            notification = {
                'jsonrpc': '2.0',
                'method': method,
                'params': params
            }
            
            message = json.dumps(notification, ensure_ascii=False)
            
            if target_connections:
                # 发送到指定连接
                send_tasks = []
                for connection_id in target_connections:
                    if connection_id in self.connections:
                        websocket = self.connections[connection_id]
                        send_tasks.append(self._send_to_connection(websocket, message, connection_id))
                
                if send_tasks:
                    await asyncio.gather(*send_tasks, return_exceptions=True)
            else:
                # 广播到所有连接
                await self.broadcast_message(message)
            
            logger.debug(f"发送通知: {method} 到 {len(target_connections or self.connections)} 个连接")
            
        except Exception as e:
            logger.error(f"发送WebSocket通知失败: {e}")
    
    def get_connection_info(self) -> Dict[str, Any]:
        """获取连接信息"""
        connections_info = {}
        
        for connection_id, websocket in self.connections.items():
            connections_info[connection_id] = {
                'remote_address': str(websocket.remote_address),
                'state': websocket.state.name if hasattr(websocket.state, 'name') else str(websocket.state),
                'closed': websocket.closed,
                'subprotocols': list(websocket.subprotocols)
            }
        
        return {
            'server_info': {
                'host': self.host,
                'port': self.port,
                'is_running': self.is_running
            },
            'connections': connections_info,
            'connection_count': len(self.connections)
        }
    
    async def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        try:
            # 检查处理器状态
            handler_healthy = self.is_running and self.websocket_server is not None
            
            # 检查服务器状态
            server_health = await self.server.health_check()
            server_healthy = server_health.get('healthy', False)
            
            # 检查连接状态
            active_connections = 0
            closed_connections = 0
            
            for websocket in self.connections.values():
                if websocket.closed:
                    closed_connections += 1
                else:
                    active_connections += 1
            
            connections_healthy = closed_connections < len(self.connections) * 0.1  # 少于10%的连接关闭
            
            overall_healthy = handler_healthy and server_healthy and connections_healthy
            
            return {
                'healthy': overall_healthy,
                'status': 'healthy' if overall_healthy else 'unhealthy',
                'checks': {
                    'handler': handler_healthy,
                    'server': server_healthy,
                    'connections': connections_healthy
                },
                'metrics': {
                    'total_connections': len(self.connections),
                    'active_connections': active_connections,
                    'closed_connections': closed_connections,
                    'server_address': f"{self.host}:{self.port}"
                },
                'server_health': server_health
            }
            
        except Exception as e:
            logger.error(f"WebSocket健康检查失败: {e}")
            return {
                'healthy': False,
                'status': 'error',
                'error': str(e)
            }
    
    async def cleanup_closed_connections(self):
        """清理已关闭的连接"""
        closed_connections = []
        
        for connection_id, websocket in list(self.connections.items()):
            if websocket.closed:
                closed_connections.append(connection_id)
        
        for connection_id in closed_connections:
            await self._cleanup_connection(connection_id)
        
        if closed_connections:
            logger.info(f"清理了 {len(closed_connections)} 个已关闭的WebSocket连接")
        
        return len(closed_connections)
    
    async def __aenter__(self):
        """异步上下文管理器入口"""
        await self.start()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.stop()


async def run_websocket_server(server: MCPServer, config: Dict[str, Any]):
    """运行WebSocket服务器"""
    handler = WebSocketHandler(server, config)
    
    try:
        await handler.start()
        
        # 保持服务器运行
        while handler.is_running:
            await asyncio.sleep(1)
            
            # 定期清理已关闭的连接
            await handler.cleanup_closed_connections()
    
    except KeyboardInterrupt:
        logger.info("收到中断信号，正在停止WebSocket服务器...")
    except Exception as e:
        logger.error(f"WebSocket服务器运行失败: {e}")
        raise
    finally:
        await handler.stop()